import os,sys
import collections
os.environ['GLOG_minloglevel'] = '2' 
 
 
import caffe
import cv2
import matplotlib.pyplot as plt
import numpy as np


def Transform_data(net,img):
   
    transform = caffe.io.Transformer( {'data':net.blobs['data'].data.shape} )
    if net.blobs['data'].data.shape[1] ==1:
        image=caffe.io.load_image(img, False)
        transform.set_transpose('data', (2,0,1) )
        transform.set_raw_scale('data', 255.0 )
    elif net.blobs['data'].data.shape[1] ==3:
        image=caffe.io.load_image(img, True)
        transform.set_transpose('data', (2,0,1) )
        transform.set_raw_scale('data', 255.0 )
        transform.set_channel_swap('data', (2,1,0) )
    
    image=transform.preprocess("data", image)
    return image
   
   
def eval_model(net_name,model_dir,pic_dir,): 
    model_list=[i for i in os.listdir(model_dir) if i.endswith("caffemodel") ]
    f=open("output.txt",'w')
    f.write("process in model : %s\n"%(net_name) ) 
     
    for index,s_model in enumerate(model_list):
        print "process in model : %s"%s_model
        picdict=collections.OrderedDict()
        label_total_dict= {"0":0,"1":0,"2":0,"3":0,"4":0,"5":0,"6":0,"7":0} 
        label_right_dict= {"0":0,"1":0,"2":0,"3":0,"4":0,"5":0,"6":0,"7":0} 
        
        totol_num=0
        right_cnt=0
        caffe.set_mode_gpu()
        model_def=net_name
        model_weight=os.path.join(model_dir,s_model)
        net=caffe.Net(model_def,model_weight,caffe.TEST)
        
        for root,subdir,images in os.walk(pic_dir):
            for image_name in images:
                totol_num+=1
                label_total_dict[ root[-1:] ]+=1
                img_path=os.path.join(root,image_name)
                net.blobs['data'].data[...]=Transform_data(net,img_path)
                output=net.forward()
                if int(root[-1:])==output['prob'][0].argmax():
                    right_cnt+=1
                    label_right_dict[root[-1:]]+=1
                    
                 
        output="(%d)weight file is %s\n"%(index+1,s_model)
        f.write(output)
        for key,value in label_total_dict.iteritems():
            output="    label %s accuracy is : %.2f\n"%(key, label_right_dict[key]*1.0/value)
            f.write(output)
        output= "total accuracy is : %.2f\n\n"%(right_cnt*1.0/totol_num)
        f.write(output)
    
    print "finish! result saved in output.txt"
    f.close()  
    

if __name__=="__main__":
    eval_model("age-test-img/lenet_age.prototxt","age-test-img/model","age-test-img/pic")
